from openpyxl import load_workbook, Workbook
from openpyxl.chart import PieChart, Reference, BarChart, Series
import os
from collections import defaultdict

def summarize_sales_data(input_file):
    # 构建输出文件名
    base_name = os.path.basename(input_file)
    output_file = os.path.join(os.path.dirname(input_file), f"summary_{base_name}")

    # 加载工作簿
    wb = load_workbook(input_file)
    ws = wb.active

    # 查找所需列的位置
    header_row = next(ws.iter_rows(min_row=1, max_row=1))
    product_col = None
    quantity_col = None
    amount_col = None
    region_col = None

    for idx, cell in enumerate(header_row, 1):
        if cell.value == "产品名称":
            product_col = idx
        elif cell.value == "销售数量":
            quantity_col = idx
        elif cell.value == "销售金额":
            amount_col = idx
        elif cell.value == "客户地区":
            region_col = idx

    if not all([product_col, quantity_col, amount_col, region_col]):
        raise ValueError("未找到必要的列：产品名称、销售数量、销售金额或客户地区")

    # 使用defaultdict进行数据汇总（按产品和地区）
    product_summary = defaultdict(lambda: {'quantity': 0, 'amount': 0})
    region_summary = defaultdict(lambda: {'quantity': 0, 'amount': 0})

    # 遍历数据行进行汇总
    for row in ws.iter_rows(min_row=2):
        product = row[product_col-1].value
        region = row[region_col-1].value
        quantity = row[quantity_col-1].value or 0
        amount = row[amount_col-1].value or 0

        if product and region:
            # 按产品汇总
            product_summary[product]['quantity'] += float(quantity)
            product_summary[product]['amount'] += float(amount)
            # 按地区汇总
            region_summary[region]['quantity'] += float(quantity)
            region_summary[region]['amount'] += float(amount)

    # 创建新工作簿并写入汇总数据
    new_wb = Workbook()
    
    # 创建产品汇总sheet
    product_ws = new_wb.active
    product_ws.title = "产品汇总"

    # 写入产品汇总数据
    headers = ["产品名称", "销售数量", "销售金额"]
    product_ws.append(headers)

    total_product_quantity = 0
    total_product_amount = 0

    for product, data in sorted(product_summary.items()):
        product_ws.append([product, data['quantity'], data['amount']])
        total_product_quantity += data['quantity']
        total_product_amount += data['amount']

    # 添加产品汇总总计行
    product_ws.append(["总计", total_product_quantity, total_product_amount])

    # 调整产品汇总表列宽
    for column in product_ws.columns:
        max_length = 0
        column_letter = column[0].column_letter
        for cell in column:
            try:
                if len(str(cell.value)) > max_length:
                    max_length = len(str(cell.value))
            except:
                pass
        adjusted_width = (max_length + 2)
        product_ws.column_dimensions[column_letter].width = adjusted_width

    # 创建地区汇总sheet
    region_ws = new_wb.create_sheet(title="地区汇总")
    
    # 写入地区汇总数据
    headers = ["客户地区", "销售数量", "销售金额"]
    region_ws.append(headers)

    total_region_quantity = 0
    total_region_amount = 0

    for region, data in sorted(region_summary.items()):
        region_ws.append([region, data['quantity'], data['amount']])
        total_region_quantity += data['quantity']
        total_region_amount += data['amount']

    # 添加地区汇总总计行
    region_ws.append(["总计", total_region_quantity, total_region_amount])

    # 调整地区汇总表列宽
    for column in region_ws.columns:
        max_length = 0
        column_letter = column[0].column_letter
        for cell in column:
            try:
                if len(str(cell.value)) > max_length:
                    max_length = len(str(cell.value))
            except:
                pass
        adjusted_width = (max_length + 2)
        region_ws.column_dimensions[column_letter].width = adjusted_width

    # 创建产品销售金额饼图
    pie = PieChart()
    pie.title = "产品销售金额分布"
    
    # 选择数据范围（不包括总计行）
    labels = Reference(product_ws, min_col=1, min_row=2, max_row=len(product_summary))
    data = Reference(product_ws, min_col=3, min_row=1, max_row=len(product_summary))
    pie.add_data(data, titles_from_data=True)
    pie.set_categories(labels)
    
    # 设置图表大小和位置
    pie.width = 15
    pie.height = 10
    product_ws.add_chart(pie, "E2")
    
    # 创建地区销售数量柱状图
    bar = BarChart()
    bar.title = "地区销售数量分布"
    bar.x_axis.title = "客户地区"
    bar.y_axis.title = "销售数量"
    
    # 选择数据范围（不包括总计行）
    categories = Reference(region_ws, min_col=1, min_row=2, max_row=len(region_summary))
    data = Reference(region_ws, min_col=2, min_row=1, max_row=len(region_summary))
    
    series = Series(data, title_from_data=True)
    bar.add_data(data, titles_from_data=True)
    bar.set_categories(categories)
    
    # 设置图表大小和位置
    bar.width = 20
    bar.height = 10
    region_ws.add_chart(bar, "E2")
    
    # 保存文件
    new_wb.save(output_file)
    return output_file

if __name__ == "__main__":
    try:
        input_file = "cleaned_演示数据.xlsx"
        if not os.path.exists(input_file):
            raise FileNotFoundError(f"文件 {input_file} 不存在")

        output_file = summarize_sales_data(input_file)
        print(f"汇总完成，已保存为: {output_file}")

    except Exception as e:
        print(f"错误: {str(e)}")